// smoothed meshes using GAC + mesh based smoothing

#include "tclap/CmdLine.h"
#include <iomanip>

#include "ioutils.h"

#include <itkCurvatureAnisotropicDiffusionImageFilter.h>
#include <itkGradientMagnitudeRecursiveGaussianImageFilter.h>
#include <itkGradientRecursiveGaussianImageFilter.h>
#include <itkLabelStatisticsImageFilter.h>
#include <itkSigmoidImageFilter.h>
#include <itkMaskNegatedImageFilter.h>
#include <itkAddImageFilter.h>
#include <itkGeodesicActiveContourLevelSetImageFilter.h>

#include "itkMorphologicalSignedDistanceTransformImageFilter.h"

#include "vtkPolyDataWriter.h"

#include "itkImageToVTKImageFilter.h"
#include "vtkWindowedSincPolyDataFilter.h"
#include "vtkImageMarchingCubes.h"


#ifdef ITK_USE_MINC2
#  include "itkMINC2ImageIOFactory.h"
#endif


#ifndef vtkFloatingPointType
#define vtkFloatingPointType float
#endif



// Aidan's trick
#include <itkSmartPointer.h>
namespace itk
{
    template <typename T>
    class Instance : public T::Pointer {
    public:
        Instance() : SmartPointer<T>( T::New() ) {}
    };
}

typedef class CmdLineType
    {
    public:
      std::string MaskIm;
      std::string RawIm;
      std::string ResultMesh;
      float sigma;
      double propScale, advecScale, curveScale;
    } CmdLineType;

void ParseCmdLine(int argc, char* argv[],
                  CmdLineType &CmdLineObj
                  )
{
    using namespace TCLAP;
    try
    {
        // Define the command line object.
        CmdLine cmd("compute absolute velocity ", ' ', "0.9");
        
        ValueArg<std::string> mArg("m","mask","mask image", true,"result","string");
        cmd.add( mArg );

        ValueArg<std::string> iArg("i","raw","raw image", true,"","string");
        cmd.add( iArg );

        ValueArg<std::string> oArg("o","output","output mesh", true,"","string");
        cmd.add( oArg );

	ValueArg<float> sArg("s","sigma","sigma for gradient filter", true,0.1,"float");
        cmd.add( sArg );
        
        ValueArg<double> propScaleArg("p","PropScaling","Geodesic Propagation term",false, 0, "float");
        cmd.add(propScaleArg);
        
        ValueArg<double> advecScaleArg("v","AdvectionScaling","Geodesic Advection term",true, 0, "float");
        cmd.add(advecScaleArg);
        
        ValueArg<double> curveScaleArg("c","CurvatureScaling","Geodesic Curvature term",true, 0, "float");
        cmd.add(curveScaleArg);
        


	cmd.parse( argc, argv );
        
        CmdLineObj.MaskIm = mArg.getValue();
	CmdLineObj.RawIm = iArg.getValue();
	CmdLineObj.ResultMesh = oArg.getValue();
	CmdLineObj.sigma = sArg.getValue();

        CmdLineObj.propScale = propScaleArg.getValue();
        CmdLineObj.advecScale = advecScaleArg.getValue();
        CmdLineObj.curveScale = curveScaleArg.getValue();

    }
    catch (ArgException &e)  // catch any exceptions
    {
        std::cerr << "error: " << e.error() << " for arg " << e.argId() << std::endl;
    }
}

/////////////////////////////////////////////////
template <class FType, class MType>
void estimateAlphaBeta(typename FType::Pointer grad, 
                       typename MType::Pointer borderZone,
		       typename MType::Pointer arteryMask,

                       float *alpha, float *beta)
{
    
    typedef typename itk::LabelStatisticsImageFilter <FType, MType > LabStatImageFilterType;   
    typename LabStatImageFilterType::Pointer labstats = LabStatImageFilterType::New ();
    labstats->SetInput(grad);
    labstats->SetLabelInput(borderZone);
    labstats->Update();
    
    float borderMean, borderSig, arteryMean;

    borderMean = labstats->GetMean(0);
    borderSig = labstats->GetSigma(0);
           
    labstats->SetLabelInput(arteryMask);
    labstats->Update();
    arteryMean = labstats->GetMean(1);
       
    std::cout<<"\t"<<borderMean << "\t" << borderSig<<"\t"<<arteryMean <<"\t"<<std::endl;
    
    float K1=0.0, K2=0.0;
    K1 = (borderMean + borderSig);
    K2 = arteryMean;
    *alpha = (K2-K1)/6;
    *beta = (K1+K2)/2; 
}
////////////////////////////////////////////////////////////////////////////////////



template <class PixType, class MaskPixType, int dim>
void doMesh(const CmdLineType &CmdLineObj)
{
  typedef typename itk::Image<PixType, dim> ImType;
  typedef typename ImType::Pointer PImType;

  typedef typename itk::Image<float, dim> FImType;
  typedef typename FImType::Pointer PFImType;

  typedef typename itk::Image<MaskPixType, dim> MaskImType;
  typedef typename MaskImType::Pointer PMaskImType;


  const double propagationScaling = CmdLineObj.propScale;
  const double advectionScaling = CmdLineObj.advecScale; 
  const double curvatureScaling = CmdLineObj.curveScale;
  const double sigma = CmdLineObj.sigma;

  PMaskImType mask = readIm<MaskImType>(CmdLineObj.MaskIm);
  PImType raw = readIm<ImType>(CmdLineObj.RawIm);

  std::cout << "Smoothing " << std::endl;
  // smooth with edge preserving filter
  PFImType gradIm;
  {
  itk::Instance< itk::CurvatureAnisotropicDiffusionImageFilter< ImType, FImType> > smoother;
  smoother->SetInput(raw);
  smoother->SetTimeStep(0.0625);
  smoother->SetNumberOfIterations(5);
  smoother->SetConductanceParameter(3.0);

  itk::Instance< itk::GradientMagnitudeRecursiveGaussianImageFilter< FImType > > GradMag;
  GradMag->SetInput(smoother->GetOutput());
  GradMag->SetSigma(sigma);
  gradIm = GradMag->GetOutput();
  gradIm->Update();
  gradIm->DisconnectPipeline();

  }

  std::cout << "DT " << std::endl;

  // create the initial level set using a distance transform
  PFImType InitialLevSet;
  {
  itk::Instance< itk::MorphologicalSignedDistanceTransformImageFilter <MaskImType, FImType> > disttrans;
//  itk::Instance< itk::MorphologicalDistanceTransformImageFilter <MaskImType, FImType> > disttrans;

  disttrans->SetInput(mask);
  disttrans->SetUseImageSpacing(true);
  disttrans->SetInsideIsPositive(true);

  InitialLevSet = disttrans->GetOutput();
  InitialLevSet->Update();
  InitialLevSet->DisconnectPipeline();
  }

//  writeIm<FImType>(InitialLevSet, "ls_init.nii.gz");
  std::cout << "borderzone " << std::endl;

  PMaskImType BorderZone, Interior;
  { 
  //The following produces a band of inputSeg image where we expect the solution to be.
  itk::Instance< itk::BinaryThresholdImageFilter<FImType, MaskImType> > bThresh;
  bThresh->SetInput(InitialLevSet);
  bThresh->SetUpperThreshold(1.5);
  bThresh->SetLowerThreshold(-1.5);
  bThresh->SetInsideValue(0);
  bThresh->SetOutsideValue(1);
  bThresh->SetReleaseDataFlag(true);
  
  BorderZone = bThresh->GetOutput();
  BorderZone->Update();
  BorderZone->DisconnectPipeline();
  
  bThresh->SetLowerThreshold(itk::NumericTraits<typename FImType::PixelType>::NonpositiveMin());
  Interior = bThresh->GetOutput();
  Interior->Update();
  Interior->DisconnectPipeline();
  }    
  //writeIm<MType>(BorderZone, CmdLineObj.TempDir + "border.nii.gz");
  //writeIm<MType>(Interior, CmdLineObj.TempDir + "inside.nii.gz");

  float alpha, beta;
    
  std::cout << "alpha/beta " << std::endl;
  estimateAlphaBeta<FImType, MaskImType>(gradIm, BorderZone, 
				       Interior, &alpha, &beta);
  
  std::cout << "Alpha = " << alpha << " Beta = " << beta << std::endl;

  PFImType FeatureImage;
  { 
  itk::Instance< itk::SigmoidImageFilter<FImType, FImType > > sigmoid;
  sigmoid->SetOutputMinimum(0.0);
  sigmoid->SetOutputMaximum(1.0);
  sigmoid->SetAlpha( alpha );
  sigmoid->SetBeta( beta );
  sigmoid->SetInput( gradIm);
  sigmoid->SetReleaseDataFlag(true);
  
  FeatureImage = sigmoid->GetOutput();
  FeatureImage->Update();
  FeatureImage->DisconnectPipeline();
  }

  PFImType narrowBandRaw;
  {       
  itk::Instance < itk::MaskNegatedImageFilter<FImType, MaskImType, FImType> > maskNeg;
  maskNeg->SetInput1(FeatureImage); //input image
  maskNeg->SetInput2(BorderZone); //mask image
  maskNeg->SetReleaseDataFlag(true);
  
  itk::Instance< itk::AddImageFilter<FImType, MaskImType, FImType> > addImage;
  addImage->SetInput1(maskNeg->GetOutput()); //input image
  addImage->SetInput2(Interior); //mask image
  addImage->SetReleaseDataFlag(true);
  
  narrowBandRaw= addImage->GetOutput();
  narrowBandRaw->SetReleaseDataFlag(true);
  narrowBandRaw->Update();
  narrowBandRaw->DisconnectPipeline();
  }

  typedef typename itk::GeodesicActiveContourLevelSetImageFilter<FImType, FImType> GeodesicActiveContourFilterType;
  typedef typename  GeodesicActiveContourFilterType::VectorImageType VectorImageType;
      
  typename VectorImageType::Pointer advecImage;
  {
  typedef typename itk::GradientRecursiveGaussianImageFilter<FImType, VectorImageType>
    GradientFilterType;
  typename GradientFilterType::Pointer gradientMapFilter = GradientFilterType::New();
  gradientMapFilter->SetInput( FeatureImage);
  gradientMapFilter->SetSigma(0.01 );
  gradientMapFilter->Update(); 
  
  advecImage = gradientMapFilter->GetOutput();
  advecImage->Update();
  advecImage->DisconnectPipeline();          
  
  //Get the negative of the above advecImage
   typedef itk::ImageRegionIterator<VectorImageType> IteratorType1;
   IteratorType1 dit( advecImage, FeatureImage->GetRequestedRegion() );
  
   for( dit.GoToBegin() ; !dit.IsAtEnd(); ++dit )
     {
     typename VectorImageType::PixelType v = dit.Get();
     for ( unsigned int j = 0; j < v.Size(); j++ )
       {
       v[j] *= -1.0L;
       }
     dit.Set( v);
     }
  }
  //FeatureImage = 0; //all memory for feature image is released  

  // With UseMinimalCurvatureOn(), tubular surfaces are smoothed along their longitudinal direction
  // while the circumferential direction is preserved.
  typename GeodesicActiveContourFilterType::Pointer geodesicActiveContour = GeodesicActiveContourFilterType::New();
  geodesicActiveContour->SetPropagationScaling( propagationScaling );
  geodesicActiveContour->SetCurvatureScaling( curvatureScaling );
  geodesicActiveContour->SetAdvectionScaling( advectionScaling );
  geodesicActiveContour->SetMaximumRMSError( 0.001 );
  geodesicActiveContour->SetNumberOfIterations(500);
  geodesicActiveContour->SetInput(InitialLevSet);//initial level set
  geodesicActiveContour->SetFeatureImage(narrowBandRaw);//feature image
  geodesicActiveContour->AutoGenerateSpeedAdvectionOff();
  geodesicActiveContour->SetAdvectionImage(advecImage); //advection image
  geodesicActiveContour->UseMinimalCurvatureOn ();
  //geodesicActiveContour->SetDerivativeSigma(0); //sets sigma for advection image 
  geodesicActiveContour->SetReleaseDataFlag(true);
  geodesicActiveContour->Update();
  
  writeIm<FImType> (geodesicActiveContour->GetOutput(), "geodesic.nii.gz");

  // now convert to a mesh and smooth
  // convert to vtk image
  typedef typename itk::ImageToVTKImageFilter<FImType> Itk2VtkType;
  typename Itk2VtkType::Pointer  m_Itk2Vtk = Itk2VtkType::New();
  
  m_Itk2Vtk->SetInput(geodesicActiveContour->GetOutput());  // m_Reader reads a binary image
  m_Itk2Vtk->Update();
  std::cout << "Image converted to VTK...." << std::endl;

  // generates iso surface and the output is vtk polydata. 
  const float objectValue = 0;
  
  vtkImageMarchingCubes *marcher = vtkImageMarchingCubes::New();
  marcher->SetInput(m_Itk2Vtk->GetOutput());
  marcher->SetValue(0, objectValue);
  marcher->Update();

  //vtkPolyData* polyData = marcher->GetOutput();
  //The following is used for smoothg of GAC mesh.
  //SetNumberOfIterations is the degree of the polynomial that is used to approximate the windowed sinc function. Set to 10 or 20.
  //FeatureEdgeSmoothingOn() helps reduce high frequency information loss that is caused by smoothing
  //SetFeatureAngle() specifies the feature angle for sharp edge identification.
  vtkWindowedSincPolyDataFilter *smoother = vtkWindowedSincPolyDataFilter::New()
;
  smoother->SetInput(marcher->GetOutput());
  smoother->SetNumberOfIterations(100);
  smoother->FeatureEdgeSmoothingOn();
  smoother->SetFeatureAngle(100);
  std::cout << "Smoothing with " << smoother->GetNumberOfIterations() << " iterations." << std::endl;

  //writes data to a file 
  vtkPolyDataWriter* writer1 = vtkPolyDataWriter::New();
  //writer1->SetFileTypeToBinary(); //sets vtk polydata to binary
  writer1->SetInput(smoother->GetOutput());
  writer1->SetFileName(CmdLineObj.ResultMesh.c_str());
  writer1->Write();

}

int main(int argc, char * argv[])
{
//    itk::MultiThreader::SetGlobalMaximumNumberOfThreads(1);
  const unsigned int dim = 3;
  CmdLineType CmdLineObj;
  ParseCmdLine(argc, argv, CmdLineObj);
#ifdef ITK_USE_MINC2
  // Enable MINC2 support
  itk::ObjectFactoryBase::RegisterFactory(itk::MINC2ImageIOFactory::New());
#endif

  doMesh<signed short, unsigned char, dim>(CmdLineObj);
  return EXIT_SUCCESS;
}
